#/usr/bin/python3
import re, hashlib 
import urllib.request, urllib.parse, urllib.error, mimetypes
import os, stat, sys, time, io
from http.cookiejar import CookieJar
import argparse

username=""
password=""
version = "0.3"

# python upload file stuff to make it all in one
# from here http://pipe.scs.fsu.edu/PostHandler/MultipartPostHandler.py
class MultipartPostHandler(urllib.request.BaseHandler):
    # needs to run first
    handler_order = urllib.request.HTTPHandler.handler_order - 10 
    
    # Controls how sequences are uncoded. 
    # If true, elements may be given multiple values by
    # assigning a sequence.
    doseq = 1

    def http_request(self, request):
        data = request.get_data()
        if data is not None and type(data) != str:
            v_files = []
            v_vars = []
            try:
                for(key, value) in list(data.items()):
                    if type(value) == io.BufferedReader:
                        v_files.append((key, value))
                    else:
                        v_vars.append((key, value))
            except TypeError:
                systype, value, traceback = sys.exc_info()
                raise TypeError("not a valid non-string sequence or mapping object %d" % traceback)

            if len(v_files) == 0:
                data = urllib.parse.urlencode(v_vars, self.doseq).encode()
            else:
                boundary, data = self.multipart_encode(v_vars, v_files)
                contenttype = 'multipart/form-data; boundary=%s' % boundary
                request.add_unredirected_header('Content-Type', contenttype)

            request.add_data(data)
        return request

    def multipart_encode(self, v_vars, files, boundary = None, buffer = None):
        if boundary is None:
            # dirty boundary
            boundary = hashlib.md5(str(time.time()).encode()).hexdigest()
        if buffer is None:
            buffer = ''
        for(key, value) in v_vars:
            buffer += '--%s\r\n' % boundary
            buffer += 'Content-Disposition: form-data; name="%s"' % key
            buffer += '\r\n\r\n' + value + '\r\n'
        for(key, fd) in files:
            file_size = os.fstat(fd.fileno())[stat.ST_SIZE]
            filename = fd.name.split('/')[-1]
            contenttype = mimetypes.guess_type(filename)[0] or \
                'application/octet-stream'
            buffer += '--%s\r\n' % boundary
            buffer += 'Content-Disposition: form-data; name="%s";'\
                ' filename="%s"\r\n' % (key, filename)
            buffer += 'Content-Type: %s\r\n' % contenttype
            fd.seek(0)
            buffer += '\r\n' + fd.read().decode('ISO-8859-1') + '\r\n'
        buffer += '--%s--\r\n\r\n' % boundary
        return boundary, buffer.encode('ISO-8859-1')

    https_request = http_request

class MalwareLuAuthError(Exception):
    def __str__(self):
        return "Wrong username or password"

class MalwareLuAlreadyHere(Exception):
    def __init__(self, md5):
        self.md5 = md5
    def __str__(self):
        return "Hash %s already in database" % self.md5
        

class MalwareLu():
    url_auth = "http://malware.lu/_auth.php"
    url_download = "http://malware.lu/_download.php"
    url_upload = "http://malware.lu/_upload.php"
    url_search = "http://malware.lu/_search.php"
    url_logout = "http://malware.lu/_logout.php"

    def __init__(self, username, password):
        self.cj = CookieJar()
        self.opener = urllib.request.build_opener(
            urllib.request.HTTPCookieProcessor(self.cj))

        if self.login(username, password) == False:
            raise MalwareLuAuthError

    def __del__(self):
        self.logout()
        
    def filetohash(self, filename):
        file_content = open(filename, "rb").read()
        sha = hashlib.md5(file_content).hexdigest()
        return sha

    def get_content_type(self, filename):
        return mimetypes.guess_type(filename)[0] or 'application/octet-stream'

    def login(self,username, password):
        params = urllib.parse.urlencode({'username': username, 'password': password})
        params = params.encode()
        f = self.opener.open(self.url_auth, params)
        data = f.read().decode()
        if not re.search("<input type=\"text\" name=\"md5\"/>", data):
            return False
        return True

    def logout(self):
        f = self.opener.open(self.url_logout)
        data = f.read()
     
    def download(self, md5):
        params = urllib.parse.urlencode({'md5': md5})
        f =  self.opener.open(self.url_download + "?%s" % params)
        params = params.encode()
        headers = f.info()
        if headers['content-type'] == "application/octet-stream":
            filename = re.findall("filename=(\S+)", 
                headers['Content-Disposition'])[0]
            fp = open(filename, 'wb')
            fp.write(f.read())
            fp.close()
            return True
        return False

    def check(self, md5):
        params = urllib.parse.urlencode({'md5': md5})
        f = self.opener.open(self.url_search + "?%s" % params)
        data = f.read().decode()
        if re.search("not found in the database", data):
            return False
        return True
        

    def upload(self, filename):
        if self.check(self.filetohash(filename)) == True:
            raise MalwareLuAlreadyHere(self.filetohash(filename))
        opener = urllib.request.build_opener(urllib.request.HTTPCookieProcessor(self.cj), 
            MultipartPostHandler)
        params = {'file': open(filename, 'rb'), }
        f = opener.open(self.url_upload, params)
        data = f.read().decode()
        if not re.search("Upload done", data):
            return False
        return True

def main():
    parser = argparse.ArgumentParser(description = "Malware.lu API tools")
    parser.add_argument('--version', action='version', 
        version="%(prog)s version " + version)
    parser.add_argument('-l', '--user', action="store", default=username)
    parser.add_argument('-p', '--passwd', action="store", default=password)
    parser.add_argument('hash', nargs="+", 
        help='filename or hash in function of action')

    group = parser.add_mutually_exclusive_group()
    group.add_argument('-c', '--check', dest="action", 
        action='store_const', const="check",
        help="Check a hash on malware.lu")
    group.add_argument('-u', '--upload', dest="action", 
        action='store_const', const="upload",
        help="Upload a file to malware.lu")
    group.add_argument('-d', '--download', dest="action", 
        action='store_const', const="download", 
        help="Download a hash from malware.lu (defaults action)")
    
    r = parser.parse_args()
    
    if len(r.user) == 0 or len(r.passwd) == 0:
        print ("Set username and password in the script or in arg")
        sys.exit(2)

    try:        
        m = MalwareLu(r.user, r.passwd)
    except MalwareLuAuthError as e:
        print (e)
        sys.exit(2)

    for arg in r.hash:
        if r.action == "upload":
            print ("Upload of %s..." % arg)
            try:
                if m.upload(arg) == True :
                    print ("Upload %s successfully" % arg)
                else:
                    print ("Upload %s failed" % arg)
            except MalwareLuAlreadyHere as e:
               print (e)

        elif r.action == "check":
            if m.check(arg) == True:
                print ("%s is in database" % arg)
            else:
                print ("%s not in database" % arg)
        elif r.action == "download" or r.action == None:
            print ("Download of %s..." % arg)
            if m.download(arg) == True:
                print ("Download %s successfully" % arg)
            else:
                print ("Download %s failed not in db or quota reach" % arg)

if __name__ == '__main__':
	main()
